#!/usr/bin/env python
# -*- coding: utf-8 -*-

from typing import List, Dict, Any, Optional, Tuple, Literal
from datetime import datetime

from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.future import select
from sqlalchemy.exc import IntegrityError
from sqlalchemy import update, delete, and_, or_

from app.core.exceptions import DatabaseError, NotFoundError
from app.core.logging import get_logger, LoggerMixin
from app.models.database import TrainingData, Project

# 创建日志记录器
logger = get_logger("training_data_repo")


class TrainingDataRepository(LoggerMixin):
    """训练数据仓库
    
    管理训练数据
    """
    
    def __init__(self, session: AsyncSession):
        """初始化训练数据仓库
        
        Args:
            session: 数据库会话
        """
        self.session = session
    
    async def create(
        self,
        project_id: str,
        data_type: Literal["ddl", "question_sql", "description"],
        content: str,
        question: Optional[str] = None,
        sql: Optional[str] = None
    ) -> Dict[str, Any]:
        """创建训练数据
        
        Args:
            project_id: 项目ID
            data_type: 数据类型
            content: 内容
            question: 问题
            sql: SQL查询
            
        Returns:
            Dict[str, Any]: 创建的训练数据字典
            
        Raises:
            DatabaseError: 当创建失败时
        """
        try:
            # 验证项目是否存在
            result = await self.session.execute(
                select(Project).filter(Project.id == project_id)
            )
            project = result.scalars().first()
            
            if not project:
                raise NotFoundError(
                    message=f"项目不存在: {project_id}",
                    resource_type="project",
                    resource_id=project_id
                )
            
            # 创建训练数据
            training_data = TrainingData(
                project_id=project_id,
                type=data_type,
                content=content,
                question=question,
                sql=sql
            )
            
            # 添加到会话
            self.session.add(training_data)
            await self.session.flush()
            
            # 创建训练数据字典
            training_data_dict = {
                "id": training_data.id,
                "project_id": project_id,
                "type": data_type,
                "content": content,
                "question": question,
                "sql": sql,
                "created_at": datetime.now().isoformat()
            }
            
            self.logger.info(f"已创建训练数据: {training_data.id}")
            return training_data_dict
            
        except NotFoundError:
            raise
        except IntegrityError as e:
            await self.session.rollback()
            error_msg = f"创建训练数据失败，可能存在完整性约束问题: {str(e)}"
            self.logger.error(error_msg)
            raise DatabaseError(error_msg)
        except Exception as e:
            await self.session.rollback()
            error_msg = f"创建训练数据失败: {str(e)}"
            self.logger.error(error_msg)
            raise DatabaseError(error_msg)
    
    async def get_by_id(self, data_id: int) -> Dict[str, Any]:
        """通过ID获取训练数据
        
        Args:
            data_id: 训练数据ID
            
        Returns:
            Dict[str, Any]: 训练数据字典
            
        Raises:
            NotFoundError: 当训练数据不存在时
        """
        try:
            # 查询训练数据
            result = await self.session.execute(
                select(TrainingData).filter(TrainingData.id == data_id)
            )
            training_data = result.scalars().first()
            
            # 检查训练数据是否存在
            if not training_data:
                raise NotFoundError(
                    message=f"训练数据不存在: {data_id}",
                    resource_type="training_data",
                    resource_id=str(data_id)
                )
            
            # 转换为字典
            return {
                "id": training_data.id,
                "project_id": training_data.project_id,
                "type": training_data.type,
                "content": training_data.content,
                "question": training_data.question,
                "sql": training_data.sql,
                "created_at": training_data.created_at.isoformat() if training_data.created_at else None
            }
            
        except NotFoundError:
            raise
        except Exception as e:
            error_msg = f"获取训练数据失败: {str(e)}"
            self.logger.error(error_msg)
            raise DatabaseError(error_msg)
    
    async def delete(self, data_id: int) -> None:
        """删除训练数据
        
        Args:
            data_id: 训练数据ID
            
        Raises:
            NotFoundError: 当训练数据不存在时
            DatabaseError: 当删除失败时
        """
        try:
            # 检查训练数据是否存在
            await self.get_by_id(data_id)
            
            # 删除训练数据
            await self.session.execute(
                delete(TrainingData).where(TrainingData.id == data_id)
            )
            await self.session.flush()
            
            self.logger.info(f"已删除训练数据: {data_id}")
            
        except NotFoundError:
            raise
        except Exception as e:
            await self.session.rollback()
            error_msg = f"删除训练数据失败: {str(e)}"
            self.logger.error(error_msg)
            raise DatabaseError(error_msg)
    
    async def list_by_project(
        self,
        project_id: str,
        data_type: Optional[str] = None,
        limit: int = 10,
        offset: int = 0
    ) -> Tuple[List[Dict[str, Any]], int]:
        """列出项目的训练数据
        
        Args:
            project_id: 项目ID
            data_type: 数据类型
            limit: 限制数量
            offset: 偏移量
            
        Returns:
            Tuple[List[Dict[str, Any]], int]: 训练数据字典列表和总数
            
        Raises:
            DatabaseError: 当查询失败时
        """
        try:
            # 构建查询条件
            conditions = [TrainingData.project_id == project_id]
            if data_type:
                conditions.append(TrainingData.type == data_type)
            
            # 查询训练数据
            result = await self.session.execute(
                select(TrainingData)
                .filter(and_(*conditions))
                .order_by(TrainingData.created_at.desc())
                .limit(limit)
                .offset(offset)
            )
            training_data_list = result.scalars().all()
            
            # 查询总数
            count_result = await self.session.execute(
                select(TrainingData.id)
                .filter(and_(*conditions))
            )
            total = len(count_result.scalars().all())
            
            # 转换为字典列表
            training_data_dicts = []
            for data in training_data_list:
                training_data_dicts.append({
                    "id": data.id,
                    "project_id": data.project_id,
                    "type": data.type,
                    "content": data.content,
                    "question": data.question,
                    "sql": data.sql,
                    "created_at": data.created_at.isoformat() if data.created_at else None
                })
            
            return training_data_dicts, total
            
        except Exception as e:
            error_msg = f"列出训练数据失败: {str(e)}"
            self.logger.error(error_msg)
            raise DatabaseError(error_msg)
    
    async def get_ddl_by_project(self, project_id: str) -> List[TrainingData]:
        """获取项目的DDL数据
        
        Args:
            project_id: 项目ID
            
        Returns:
            List[TrainingData]: DDL数据列表
            
        Raises:
            DatabaseError: 当查询失败时
        """
        try:
            # 查询DDL数据
            result = await self.session.execute(
                select(TrainingData)
                .filter(
                    and_(
                        TrainingData.project_id == project_id,
                        TrainingData.type == "ddl"
                    )
                )
                .order_by(TrainingData.created_at.desc())
            )
            ddl_list = result.scalars().all()
            
            return list(ddl_list)
            
        except Exception as e:
            error_msg = f"获取DDL数据失败: {str(e)}"
            self.logger.error(error_msg)
            raise DatabaseError(error_msg)
    
    async def get_question_sql_by_ids(self, data_ids: List[int]) -> List[TrainingData]:
        """通过ID列表获取问题-SQL数据
        
        Args:
            data_ids: 数据ID列表
            
        Returns:
            List[TrainingData]: 问题-SQL数据列表
            
        Raises:
            DatabaseError: 当查询失败时
        """
        if not data_ids:
            return []
        
        try:
            # 查询问题-SQL数据
            result = await self.session.execute(
                select(TrainingData)
                .filter(
                    and_(
                        TrainingData.id.in_(data_ids),
                        TrainingData.type == "question_sql"
                    )
                )
            )
            question_sql_list = result.scalars().all()
            
            return list(question_sql_list)
            
        except Exception as e:
            error_msg = f"获取问题-SQL数据失败: {str(e)}"
            self.logger.error(error_msg)
            raise DatabaseError(error_msg)
            
    async def get_ddl_by_ids(self, data_ids: List[int]) -> List[TrainingData]:
        """通过ID列表获取DDL数据
        
        Args:
            data_ids: 数据ID列表
            
        Returns:
            List[TrainingData]: DDL数据列表
            
        Raises:
            DatabaseError: 当查询失败时
        """
        if not data_ids:
            return []
        
        try:
            # 查询DDL数据
            result = await self.session.execute(
                select(TrainingData)
                .filter(
                    and_(
                        TrainingData.id.in_(data_ids),
                        TrainingData.type == "ddl"
                    )
                )
            )
            ddl_list = result.scalars().all()
            
            return list(ddl_list)
            
        except Exception as e:
            error_msg = f"获取DDL数据失败: {str(e)}"
            self.logger.error(error_msg)
            raise DatabaseError(error_msg)
    
    async def get_descriptions_by_ids(self, data_ids: List[int]) -> List[TrainingData]:
        """通过ID列表获取描述数据
        
        Args:
            data_ids: 数据ID列表
            
        Returns:
            List[TrainingData]: 描述数据列表
            
        Raises:
            DatabaseError: 当查询失败时
        """
        if not data_ids:
            return []
        
        try:
            # 查询描述数据
            result = await self.session.execute(
                select(TrainingData)
                .filter(
                    and_(
                        TrainingData.id.in_(data_ids),
                        TrainingData.type == "description"
                    )
                )
            )
            description_list = result.scalars().all()
            
            return list(description_list)
            
        except Exception as e:
            error_msg = f"获取描述数据失败: {str(e)}"
            self.logger.error(error_msg)
            raise DatabaseError(error_msg)
    
    async def search_by_content(
        self,
        project_id: str,
        query: str,
        data_type: Optional[str] = None,
        limit: int = 10
    ) -> List[TrainingData]:
        """通过内容搜索训练数据
        
        Args:
            project_id: 项目ID
            query: 搜索查询
            data_type: 数据类型
            limit: 限制数量
            
        Returns:
            List[TrainingData]: 训练数据列表
            
        Raises:
            DatabaseError: 当搜索失败时
        """
        try:
            # 构建查询条件
            conditions = [TrainingData.project_id == project_id]
            
            # 添加数据类型条件
            if data_type:
                conditions.append(TrainingData.type == data_type)
            
            # 添加内容搜索条件
            search_condition = or_(
                TrainingData.content.like(f"%{query}%"),
                TrainingData.question.like(f"%{query}%") if query else False
            )
            conditions.append(search_condition)
            
            # 执行查询
            result = await self.session.execute(
                select(TrainingData)
                .filter(and_(*conditions))
                .limit(limit)
            )
            training_data_list = result.scalars().all()
            
            return list(training_data_list)
            
        except Exception as e:
            error_msg = f"搜索训练数据失败: {str(e)}"
            self.logger.error(error_msg)
            raise DatabaseError(error_msg)